SageMaker HyperPod をユーザー所有の VPC にホストする

#SageMaker HyperPod
#AWS
たかくに
2024.12.14
こんにちは！AWS 事業本部コンサルティング部のたかくに（@takakuni_）です。
Amazon SageMaker HyperPod 使っていますか？
今回は SageMaker HyperPod クラスターをユーザーが所有している VPC に所属させます。
図にすると以下の構成になります。
 ユーザー管理の VPC 所属させるということユースケースを考えてみましょう。以下が考えられますね。
まず第一想起するのは、オーケストレータに EKS を使いたいケースでしょうか。Slurm でもネットワーク経路を制御できるものにしたい、固定の IP アドレスからデータを取得したいケースが考えられます。
SageMaker HyperPod クラスターをユーザーが所有している VPC に所属させるということは、ユーザーは VPC を意識する必要があります。
我ながら、進次郎構文がクリティカルに炸裂していますね。
説明に移ります。インスタンスグループは構成図の通り、ユーザーが指定したサブネット内に ENI を生やす形で VPC に所属します。つまり、インターネットや S3 向けの通信は指定したサブネットのルートテーブルに従い通信し、通信制御はセキュリティグループ、Network ACL で制御します。
伝えたかったことは、これだけではありません。
ネットワークレイテンシーを最大限に高めるには EFA 対応のインスタンスタイプを選んだり、EFA ヘルスチェック用にセキュリティグループに更新したり、インスタンス群を 1AZ に集約させる必要があります。[1]
詳しい情報や最新情報は以下をご覧ください。
https://docs.aws.amazon.com/sagemaker/latest/dg/sagemaker-hyperpod-prerequisites.html#sagemaker-hyperpod-prerequisites-optional-vpc
 やってみるそれでは今回はユーザー所有の VPC に SageMaker HyperPod Cluster を所属させてみたいと思います。
再掲になりますが構成図は以下になります。
リソースの作成は HashiCorp Terraform で作成しました。利用したコードの全文は以下に格納しています。
https://github.com/takakuni-classmethod/genai-blog/tree/main/sagemaker_hyperpod_vpc
 IAMVPC に所属させるため IAM には AmazonSageMakerClusterInstanceRolePolicy に加えて、 以下の権限を追加させる必要があります。
{
    "Effect": "Allow",
    "Action": [
        "ec2:CreateNetworkInterface",
        "ec2:CreateNetworkInterfacePermission",
        "ec2:DeleteNetworkInterface",
        "ec2:DeleteNetworkInterfacePermission",
        "ec2:DescribeNetworkInterfaces",
        "ec2:DescribeVpcs",
        "ec2:DescribeDhcpOptions",
        "ec2:DescribeSubnets",
        "ec2:DescribeSecurityGroups",
        "ec2:DetachNetworkInterface"
    ],
    "Resource": "*"
}
{
    "Effect": "Allow",
    "Action": "ec2:CreateTags",
    "Resource": [
        "arn:aws:ec2:*:*:network-interface/*"
    ]
}
https://docs.aws.amazon.com/sagemaker/latest/dg/sagemaker-hyperpod-prerequisites-iam.html#sagemaker-hyperpod-prerequisites-iam-role-for-hyperpod
 VPC ConfigSageMaker HyperPod の CreateCluster API を見てみましょう。
VPC が設定可能な箇所が VpcConfig と OverrideVpcConfig の 2 種類あります。
{
   "ClusterName": "string",
   "InstanceGroups": [
      {
         "ExecutionRole": "string",
         "InstanceCount": number,
         "InstanceGroupName": "string",
         "InstanceStorageConfigs": [
            { ... }
         ],
         "InstanceType": "string",
         "LifeCycleConfig": {
            "OnCreate": "string",
            "SourceS3Uri": "string"
         },
         "OnStartDeepHealthChecks": [ "string" ],
         "OverrideVpcConfig": {
            "SecurityGroupIds": [ "string" ],
            "Subnets": [ "string" ]
         },
         "ThreadsPerCore": number,
         "TrainingPlanArn": "string"
      }
   ],
   "NodeRecovery": "string",
   "Orchestrator": {
      "Eks": {
         "ClusterArn": "string"
      }
   },
   "Tags": [
      {
         "Key": "string",
         "Value": "string"
      }
   ],
   "VpcConfig": {
      "SecurityGroupIds": [ "string" ],
      "Subnets": [ "string" ]
   }
}
https://docs.aws.amazon.com/sagemaker/latest/APIReference/API_CreateCluster.html
OverrideVpcConfig はインスタンスグループ単位で VPC 設定を上書きできる設定値です。OverrideVpcConfig を設定しない場合、クラスターで設定した VpcConfig が適用されるイメージです。
つまり、次のようにインスタンスグループ別にサブネットやセキュリティを使い分けることができたりします。
セキュリティグループも分けられるということは、マウントしたい Lustre ファイルシステムもうまく制御できそうです。
サブネットを分ける例としては、各インスタンスグループ別に NAT Gateway の帯域を占有したいケースが挙げられますかね。（リッチですね。）
 Security Group今回、 EFA を使うわけではないのですが、例に倣って次のルールを設定しました。
 インバウンド

タイプ
プロトコル
ポート
ソース


すべてのトラフィック
すべて
すべて
自身のセキュリティグループ ID

 アウトバウンド

タイプ
プロトコル
ポート
ソース


すべてのトラフィック
すべて
すべて
0.0.0.0/0

 疎通確認最後に疎通確認をして終わります。マネジメントコンソールから NAT Gateway の EIP は 18.176. から始まっていることがわかります。
SSM で接続してパブリック IP をチェックしてみます。
NAT Gateway の IP を利用して、インターネットへ抜けていることがわかりますね。
[cloudshell-user@ip-10-132-76-34 ~]$ aws ssm start-session \
>     --target sagemaker-cluster:px67k0kfjbgr_controller-group-i-0f8d40cb0d31152ca \
>     --region ap-northeast-1

Starting session with SessionId: takakuni-qbf9zgjrpha2h9v8eleql5zv2u
# curl http://checkip.amazonaws.com/
18.176.XXX.XXX
#
 まとめ以上、「SageMaker HyperPod をユーザー所有の VPC にホストする」でした。
一昔前に以下のブログを書いたのを思い出し、SageMaker HyperPod の場合どうなるのか気になったのでブログにしてみました。
https://devio2024-front-preview.developers.io/articles/sagemaker-network-vpc-architecture-2022-04/
もし、ネットワーク周りをユーザー管理の VPC に所属させたくなった場合に見返そうと思います。
このブログがどなたかの参考になれば幸いです。
AWS 事業本部コンサルティング部のたかくに（@takakuni_）でした！

脚注
もちろん可用性は低下します。 ↩︎